package org.znxs.znmanus.rag.a_etl;

import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.markdown.MarkdownDocumentReader;
import org.springframework.ai.reader.markdown.config.MarkdownDocumentReaderConfig;
import org.springframework.core.io.Resource;
import org.springframework.core.io.support.ResourcePatternResolver;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

/**
 * 文档加载器 从本地加载MarkDown文件到内存中
 */
@Component
@Slf4j
public class MarkDownLoader {

    private final ResourcePatternResolver resourcePatternResolver;


    public MarkDownLoader(ResourcePatternResolver resourcePatternResolver) {
        this.resourcePatternResolver = resourcePatternResolver;
    }


    /**
     * 加载文档
     *
     * @param fileFolder 文件所在文件夹
     * @return
     */
    public List<Document> loadMarkDowns(String fileFolder) {
        ArrayList<Document> documents = new ArrayList<>();

        try {
            Resource[] resources = resourcePatternResolver.getResources(fileFolder);
            for (Resource resource : resources) {
                // 获取文件名
                String filename = resource.getFilename();
                // 获取元信息 直接将文件名作为元数据
                // 适用spring ai 提供的md 文档读取进行构建配置
                MarkdownDocumentReaderConfig config = MarkdownDocumentReaderConfig.builder()
                        .withHorizontalRuleCreateDocument(true) // 水平线是否生成新文档
                        .withIncludeCodeBlock(false)  // 是否包含代码块
                        .withIncludeBlockquote(false) // 是否包含引用块
                        // 将文件名称作为元信息
                        .withAdditionalMetadata("filename", filename)
                        .build();
                // 进行读取
                MarkdownDocumentReader markdownDocumentReader = new MarkdownDocumentReader(resource, config);
                // 把文档保存到集合中
                documents.addAll(markdownDocumentReader.get());
            }
            return documents;

        } catch (IOException e) {
            throw new RuntimeException("文档加载出现问题" + e);
        }


    }

}